#coding:utf-8
import sys
import os
import re
import time
import requests
import threading
import ipaddress
import urllib3
import psutil
from PyQt5.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QWidget,
                            QLabel, QTextEdit, QPushButton, QProgressBar, 
                            QHBoxLayout, QFileDialog, QSpinBox, QGroupBox, 
                            QMessageBox, QSplitter, QComboBox, QLineEdit)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
from urllib.parse import urljoin, urlparse
from bs4 import BeautifulSoup
import concurrent.futures
from itertools import product

# 禁用SSL警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

class SystemMonitor:
    """系统资源监控类"""
    @staticmethod
    def get_cpu_usage():
        try:
            return psutil.cpu_percent(interval=0.1)
        except:
            return "N/A"

    @staticmethod
    def get_memory_usage():
        try:
            mem = psutil.virtual_memory()
            return f"{mem.percent}% ({round(mem.used/1024/1024)}MB/{round(mem.total/1024/1024)}MB)"
        except:
            return "N/A"

    @staticmethod 
    def get_disk_usage():
        try:
            disk = psutil.disk_usage('/')
            return f"{disk.percent}% ({round(disk.used/1024/1024/1024, 1)}GB/{round(disk.total/1024/1024/1024, 1)}GB)"
        except:
            return "N/A"

class JSFinder:
    @staticmethod
    def extract_urls(content):
        pattern_raw = r"""
      (?:"|')
      (
        ((?:[a-zA-Z]{1,10}://|//)
        [^"'/]{1,}\.
        [a-zA-Z]{2,}[^"']{0,})
        |
        ((?:/|\.\./|\./)
        [^"'><,;| *()(%$^/\\\[\]]
        [^"'><,;|()]{1,})
        |
        ([a-zA-Z0-9_\-/]{1,}/
        [a-zA-Z0-9_\-/]{1,}
        \.(?:[a-zA-Z]{1,4}|action)
        (?:[\?|/][^"|']{0,}|))
        |
        ([a-zA-Z0-9_\-]{1,}
        \.(?:php|asp|aspx|jsp|json|action|html|js|txt|xml)
        (?:\?[^"|']{0,}|))
      )
      (?:"|')
        """
        pattern = re.compile(pattern_raw, re.VERBOSE)
        return [match.group().strip('"').strip("'") for match in re.finditer(pattern, str(content))]

    @staticmethod
    def get_html(url):
        try:
            response = requests.get(
                url, 
                headers={'User-Agent': 'Mozilla/5.0'}, 
                timeout=8, 
                verify=False,
                allow_redirects=True
            )
            return response.content.decode("utf-8", "ignore")
        except Exception as e:
            print(f"获取HTML失败: {str(e)}")
            return None

    @staticmethod
    def find_urls(target_url):
        if not target_url.startswith(('http://', 'https://')):
            target_url = f'http://{target_url}'
            
        html = JSFinder.get_html(target_url)
        if not html:
            return []
        
        soup = BeautifulSoup(html, "html.parser")
        scripts = soup.findAll("script")
        script_contents = {}
        inline_script = ""
        
        for script in scripts:
            src = script.get("src")
            if not src:
                inline_script += script.get_text() + "\n"
            else:
                full_url = urljoin(target_url, src)
                script_content = JSFinder.get_html(full_url)
                if script_content:
                    script_contents[full_url] = script_content
        
        script_contents[target_url] = inline_script
        all_urls = []
        
        for script_url in script_contents:
            urls = JSFinder.extract_urls(script_contents[script_url])
            if urls:
                for url in urls:
                    full_url = urljoin(script_url, url)
                    all_urls.append(full_url)
        
        return sorted(set(all_urls))

class PathProcessor:
    @staticmethod
    def extract_components(urls):
        components = set()
        for url in urls:
            parsed = urlparse(url)
            path = parsed.path
            if not path or path == '/':
                continue
                
            parts = [p for p in path.split('/') if p]
            for i in range(1, len(parts)+1):
                components.add('/'.join(parts[:i]))
        
        return sorted(components, key=lambda x: (x.count('/'), x))

class TargetValidator:
    @staticmethod
    def is_valid_target(target):
        """验证目标URL是否有效"""
        try:
            try:
                ipaddress.ip_address(target)
                return True
            except ValueError:
                pass
                
            parsed = urlparse(target)
            if not parsed.scheme:
                parsed = urlparse(f'http://{target}')
                
            if not parsed.netloc:
                return False
                
            if parsed.scheme not in ('http', 'https'):
                return False
                
            return True
        except:
            return False

    @staticmethod
    def is_alive(target):
        """检查目标是否存活"""
        try:
            if not target.startswith(('http://', 'https://')):
                target = f'http://{target}'
                
            response = requests.head(
                target,
                timeout=10,
                allow_redirects=True,
                headers={'User-Agent': 'Mozilla/5.0'},
                verify=False
            )
            return response.status_code < 400
        except:
            return False

class BruteWorker(QThread):
    update_progress = pyqtSignal(int, int)
    found_path = pyqtSignal(str, str, str)
    error_path = pyqtSignal(str, str)
    finished = pyqtSignal()
    status_update = pyqtSignal(str)
    paused = pyqtSignal(bool)
    thread_status = pyqtSignal(int)
    current_url = pyqtSignal(str)
    config_changed = pyqtSignal(int, int, str, int, int)
    depth_decreased = pyqtSignal()
    config_updated = pyqtSignal(str)

    def __init__(self, base_url, components, max_depth=3, max_threads=10, 
                 base_url_path="", request_delay=100, retry_count=1):
        super().__init__()
        # 解析基础URL，确保正确提取协议和域名
        parsed_url = urlparse(base_url)
        if not parsed_url.scheme:
            parsed_url = urlparse(f'http://{base_url}')
        self.original_base_url = f"{parsed_url.scheme}://{parsed_url.netloc}".rstrip('/')
        
        self.base_url_path = base_url_path.strip('/')
        self.base_url = f"{self.original_base_url}/{self.base_url_path}" if self.base_url_path else self.original_base_url
        self.components = components
        self.max_depth = max_depth
        self.max_threads = min(max_threads, 20)
        self.running = True
        self.pause_flag = False
        self.visited = set()
        self.valid_paths = set()
        self.current_depth = 1
        self.total_paths = 0
        self.completed_paths = 0
        self.request_delay = request_delay
        self.retry_count = retry_count
        self.lock = threading.Lock()
        self.futures = set()
        self.executor = None
        self.paths_to_process = []
        
        self.config_changed.connect(self.update_config)

    def pause(self):
        """暂停扫描"""
        with self.lock:
            self.pause_flag = True
        self.paused.emit(True)
        self.status_update.emit("已暂停")

    def resume(self):
        """继续扫描"""
        with self.lock:
            self.pause_flag = False
        self.paused.emit(False)
        self.status_update.emit("继续扫描")

    def update_config(self, new_depth, new_threads, new_base_path, new_delay, new_retry):
        """实时更新扫描配置（线程安全实现）"""
        with self.lock:
            old_depth = self.max_depth
            old_threads = self.max_threads
            old_base_path = self.base_url_path
            old_delay = self.request_delay
            old_retry = self.retry_count
            
            depth_changed = new_depth != old_depth
            depth_decreased = new_depth < old_depth
            
            # 更新基础配置
            self.max_depth = new_depth
            self.max_threads = min(new_threads, 20)
            self.base_url_path = new_base_path.strip('/')
            self.request_delay = new_delay
            self.retry_count = new_retry
            
            # 更新基础URL
            self.base_url = f"{self.original_base_url}/{self.base_url_path}" if self.base_url_path else self.original_base_url
            
            # 重新计算总路径数
            self.total_paths = 0
            for i in range(1, self.max_depth+1):
                self.total_paths += min(len(self.components)** i, 10000)
            
            # 线程池更新逻辑
            if self.executor and old_threads != self.max_threads:
                old_executor = self.executor
                self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.max_threads)
                self.config_updated.emit(f"线程数量已更新为: {self.max_threads}")
                
                # 安全关闭旧线程池
                def shutdown_old_executor(executor):
                    executor.shutdown(wait=False, cancel_futures=False)
                threading.Thread(target=shutdown_old_executor, args=(old_executor,), daemon=True).start()
            
            # 基础路径变更处理
            if old_base_path != self.base_url_path:
                old_base_url = f"{self.original_base_url}/{old_base_path}" if old_base_path else self.original_base_url
                
                # 过滤与旧路径相关的记录
                filtered_visited = set()
                for url in self.visited:
                    if not url.startswith(old_base_url):
                        filtered_visited.add(url)
                self.visited = filtered_visited
                
                filtered_valid = set()
                for url in self.valid_paths:
                    if not url.startswith(old_base_url):
                        filtered_valid.add(url)
                self.valid_paths = filtered_valid
                
                self.config_updated.emit(f"基础URL已更新为: {self.base_url_path or '无'}\n已重置相关访问记录")
        
        # 处理深度变更
        config_msg = []
        if depth_changed:
            config_msg.append(f"深度: {self.max_depth}")
            if depth_decreased:
                # 深度减小，过滤超出新深度的路径
                with self.lock:
                    self.paths_to_process = [(p, d) for p, d in self.paths_to_process if d <= self.max_depth]
                self.depth_decreased.emit()
            else:
                # 深度增加，生成新增深度的路径
                self.generate_remaining_paths()
                self.config_updated.emit(f"扫描深度已更新为: {self.max_depth}")
        
        # 处理其他配置变更
        if old_threads != self.max_threads:
            config_msg.append(f"线程: {self.max_threads}")
        if old_delay != self.request_delay:
            config_msg.append(f"延迟: {self.request_delay}ms")
        if old_retry != self.retry_count:
            config_msg.append(f"重试: {self.retry_count}次")
        if old_base_path != self.base_url_path:
            config_msg.append(f"基础路径: {self.base_url_path or '无'}")
            
        if config_msg:
            self.status_update.emit(f"配置已更新 - {', '.join(config_msg)}")

    def generate_remaining_paths(self):
        with self.lock:
            current_max_depth = max([d for _, d in self.paths_to_process], default=0) if self.paths_to_process else 0
            existing_paths = {path for path, _ in self.paths_to_process}
        
        # 只生成新增深度的路径
        for depth in range(current_max_depth + 1, self.max_depth + 1):
            count = 0
            max_per_depth = 10000
            for parts in product(self.components, repeat=depth):
                count += 1
                if count > max_per_depth:
                    break
                path = '/'.join(parts)
                if path not in existing_paths:
                    self.paths_to_process.append((path, depth))
                    existing_paths.add(path)

    def run(self):
        self.status_update.emit("开始扫描...")
        self.generate_all_paths()
        
        with self.lock:
            self.total_paths = len(self.paths_to_process)
        self.completed_paths = 0
        
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.max_threads)
        
        try:
            # 主循环：处理任务队列和已提交任务
            while (self.paths_to_process or self.futures) and self.running:
                # 暂停状态处理
                while self.pause_flag and self.running:
                    time.sleep(0.1)
                    self.paused.emit(True)
                
                if not self.running:
                    break
                
                # 获取当前有效的线程池
                current_executor = None
                with self.lock:
                    if self.executor:
                        current_executor = self.executor
                
                if not current_executor:
                    break
                
                # 控制并发任务数量，避免过多堆积
                with self.lock:
                    current_futures = len(self.futures)
                
                # 当并发任务数小于线程数的2倍时，添加新任务
                while (current_futures < self.max_threads * 2 
                       and self.paths_to_process 
                       and self.running 
                       and not self.pause_flag):
                    with self.lock:
                        path, depth = self.paths_to_process.pop(0)
                        self.current_depth = depth
                    
                    # 提交任务到当前线程池
                    future = current_executor.submit(self.check_path, path)
                    with self.lock:
                        self.futures.add(future)
                    future.add_done_callback(self.handle_future_complete)
                    self.update_thread_count()
                    
                    with self.lock:
                        current_futures = len(self.futures)
                
                # 短暂休眠，降低CPU占用
                time.sleep(0.05)
            
            # 等待所有任务完成
            while self.running and len(self.futures) > 0:
                time.sleep(0.1)
                
        finally:
            with self.lock:
                if self.executor:
                    self.executor.shutdown(wait=False, cancel_futures=True)
                    self.executor = None
        
        self.finished.emit()

    def generate_all_paths(self):
        self.paths_to_process = []
        seen_paths = set()  # 用于去重
        
        for depth in range(1, self.max_depth+1):
            count = 0
            max_per_depth = 10000  # 限制单深度路径数量
            for parts in product(self.components, repeat=depth):
                count += 1
                if count > max_per_depth:
                    break
                path = '/'.join(parts)
                if path not in seen_paths:  # 确保路径唯一
                    seen_paths.add(path)
                    self.paths_to_process.append((path, depth))

    def check_path(self, path):
        with self.lock:
            current_base_url = self.base_url
            current_delay = self.request_delay / 1000
            current_retry = self.retry_count
        
        full_url = urljoin(f"{current_base_url}/", path)
        last_error = None
        
        # 检查是否已访问，避免重复扫描
        with self.lock:
            if full_url in self.visited:
                return (False, full_url, "已访问过")
            self.visited.add(full_url)
        
        for _ in range(current_retry + 1):
            # 检查是否已停止
            if not self.running:
                return (False, full_url, "已停止")
                
            # 检查暂停状态
            if self.pause_flag:
                while self.pause_flag and self.running:
                    time.sleep(0.1)  # 暂停时持续等待
                if not self.running:
                    return (False, full_url, "已停止")
            
            try:
                # 根据选择的方法发送请求
                response = requests.head(
                    full_url, 
                    timeout=10, 
                    allow_redirects=True,
                    headers={'User-Agent': 'Mozilla/5.0'},
                    verify=False
                )
                status = response.status_code
                return (True, full_url, f"状态码: {status}")
            except Exception as e:
                last_error = f"错误: {str(e)}"
            
            # 仅在非最后一次重试时添加延迟
            if _ < current_retry:
                time.sleep(current_delay)
        
        return (False, full_url, last_error)

    def handle_future_complete(self, future):
        try:
            if not self.running:
                return
                
            result = future.result()
            self.process_result(*result)
        except Exception as e:
            print(f"任务处理错误: {str(e)}")
            
        with self.lock:
            if future in self.futures:
                self.futures.remove(future)
        self.update_thread_count()

    def process_result(self, success, url, info):
        """处理检查结果"""
        with self.lock:
            self.completed_paths += 1
            if self.total_paths > 0:
                progress = int((self.completed_paths / self.total_paths) * 100)
            else:
                progress = 0
        self.update_progress.emit(progress, self.completed_paths)
        self.current_url.emit(url)
        
        if success:
            with self.lock:
                self.valid_paths.add(url)
            self.found_path.emit(url, info, str(self.current_depth))
        else:
            self.error_path.emit(url, info)

    def update_thread_count(self):
        """更新线程计数"""
        with self.lock:
            count = len(self.futures)
        self.thread_status.emit(count)

    def stop(self):
        with self.lock:
            self.running = False
            self.paths_to_process = []
            self.pause_flag = False
            
            for future in self.futures:
                if not future.done():
                    future.cancel()
            self.futures.clear()
        
        self.status_update.emit("正在停止...")

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("信安西部-明镜高悬实验室 智能路由与API扫描系统（v2.0）")
        self.setup_ui()
        self.alive_targets = []
        self.current_target_index = 0
        self.results_dir = "scan_results"
        self.worker = None
        self.success_paths = []
        self.current_scan_file = None
        self.last_depth = 3
        self.last_threads = 5
        self.last_basepath = ""
        self.last_delay = 100
        self.last_retry = 1
        self.is_stopping = False
        self.batch_scanning = False
        self.custom_components = None
        self.is_paused = False
        os.makedirs(self.results_dir, exist_ok=True)
        
        self.monitor_timer = QTimer()
        self.monitor_timer.timeout.connect(self.update_monitor)
        self.monitor_timer.start(1000)
        
        self.config_timer = QTimer()
        self.config_timer.setSingleShot(True)
        self.config_timer.timeout.connect(self.apply_config_changes)

    def setup_ui(self):
        central_widget = QWidget()
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(5, 5, 5, 5)
        main_layout.setSpacing(5)

        monitor_group = QGroupBox("系统资源监控")
        monitor_layout = QHBoxLayout()
        
        self.cpu_label = QLabel("CPU: 0%")
        self.mem_label = QLabel("内存: 0/0 MB")
        self.disk_label = QLabel("磁盘: 0/0 GB")
        self.thread_label = QLabel("线程: 0")
        
        monitor_layout.addWidget(self.cpu_label)
        monitor_layout.addWidget(self.mem_label)
        monitor_layout.addWidget(self.disk_label)
        monitor_layout.addWidget(self.thread_label)
        monitor_group.setLayout(monitor_layout)
        main_layout.addWidget(monitor_group)

        # 目标和配置区域
        config_splitter = QSplitter(Qt.Horizontal)
        
        # 目标URL区域
        url_group = QGroupBox("目标设置")
        url_layout = QVBoxLayout()
        
        input_layout = QHBoxLayout()
        input_layout.addWidget(QLabel("目标URL:"))
        self.url_input = QLineEdit()
        self.url_input.setPlaceholderText("http://example.com/ 或 IP地址")
        input_layout.addWidget(self.url_input, 1)
        
        self.batch_btn = QPushButton("批量导入")
        self.batch_btn.clicked.connect(self.batch_import)
        input_layout.addWidget(self.batch_btn)
        
        url_layout.addLayout(input_layout)
        
        # 批量扫描状态显示
        self.batch_status = QLabel("批量扫描: 未启用")
        url_layout.addWidget(self.batch_status)
        
        # 基础URL区域
        baseurl_layout = QHBoxLayout()
        baseurl_layout.addWidget(QLabel("基础路径:"))
        self.baseurl_input = QLineEdit()
        self.baseurl_input.setPlaceholderText("例如: admin (将扫描 http://目标/admin/...)")
        self.baseurl_input.textChanged.connect(self.on_config_changed)
        baseurl_layout.addWidget(self.baseurl_input, 1)
        url_layout.addLayout(baseurl_layout)
        
        # 自定义字典导入区域
        dict_layout = QHBoxLayout()
        dict_layout.addWidget(QLabel("路径字典:"))
        self.dict_path_edit = QLineEdit()
        self.dict_path_edit.setPlaceholderText("自定义路径字典文件路径")
        dict_layout.addWidget(self.dict_path_edit, 1)
        
        self.import_dict_btn = QPushButton("导入字典")
        self.import_dict_btn.clicked.connect(self.import_custom_dict)
        dict_layout.addWidget(self.import_dict_btn)
        url_layout.addLayout(dict_layout)
        
        # 字典预览区域
        self.dict_preview = QTextEdit()
        self.dict_preview.setReadOnly(True)
        self.dict_preview.setMaximumHeight(80)
        url_layout.addWidget(QLabel("字典预览 (前10行):"))
        url_layout.addWidget(self.dict_preview)
        
        # 扫描控制按钮
        btn_layout = QHBoxLayout()
        self.fetch_btn = QPushButton("提取路径")
        self.fetch_btn.clicked.connect(self.fetch_urls)
        btn_layout.addWidget(self.fetch_btn)
        
        self.start_btn = QPushButton("开始扫描")
        self.start_btn.clicked.connect(self.start_scan)
        btn_layout.addWidget(self.start_btn)
        
        self.pause_btn = QPushButton("暂停")
        self.pause_btn.clicked.connect(self.toggle_pause)
        self.pause_btn.setEnabled(False)
        btn_layout.addWidget(self.pause_btn)
        
        self.stop_btn = QPushButton("停止")
        self.stop_btn.clicked.connect(self.stop_scan)
        self.stop_btn.setEnabled(False)
        btn_layout.addWidget(self.stop_btn)
        
        self.export_btn = QPushButton("导出结果")
        self.export_btn.clicked.connect(self.export_results)
        self.export_btn.setEnabled(False)
        btn_layout.addWidget(self.export_btn)
        
        url_layout.addLayout(btn_layout)
        url_group.setLayout(url_layout)
        config_splitter.addWidget(url_group)
        
        # 扫描配置区域
        scan_group = QGroupBox("扫描配置")
        scan_layout = QVBoxLayout()
        
        depth_layout = QHBoxLayout()
        depth_layout.addWidget(QLabel("最大深度:"))
        self.depth_input = QSpinBox()
        self.depth_input.setRange(1, 10)
        self.depth_input.setValue(3)
        self.depth_input.valueChanged.connect(self.on_config_changed)
        depth_layout.addWidget(self.depth_input)
        scan_layout.addLayout(depth_layout)
        
        thread_layout = QHBoxLayout()
        thread_layout.addWidget(QLabel("线程数:"))
        self.thread_input = QSpinBox()
        self.thread_input.setRange(1, 20)
        self.thread_input.setValue(5)
        self.thread_input.valueChanged.connect(self.on_config_changed)
        thread_layout.addWidget(self.thread_input)
        scan_layout.addLayout(thread_layout)
        
        # 请求延迟配置
        delay_layout = QHBoxLayout()
        delay_layout.addWidget(QLabel("请求延迟(ms):"))
        self.delay_input = QSpinBox()
        self.delay_input.setRange(0, 2000)
        self.delay_input.setValue(100)
        self.delay_input.valueChanged.connect(self.on_config_changed)
        delay_layout.addWidget(self.delay_input)
        scan_layout.addLayout(delay_layout)
        
        # 重试次数配置
        retry_layout = QHBoxLayout()
        retry_layout.addWidget(QLabel("重试次数:"))
        self.retry_input = QSpinBox()
        self.retry_input.setRange(0, 5)
        self.retry_input.setValue(1)
        self.retry_input.valueChanged.connect(self.on_config_changed)
        retry_layout.addWidget(self.retry_input)
        scan_layout.addLayout(retry_layout)
        
        # 状态码过滤配置
        status_code_layout = QHBoxLayout()
        status_code_layout.addWidget(QLabel("关注的状态码:"))
        self.status_codes_edit = QLineEdit()
        self.status_codes_edit.setPlaceholderText("用逗号分隔，如: 200,403,301")
        self.status_codes_edit.setText("200,403,301,302")  # 默认值
        status_code_layout.addWidget(self.status_codes_edit)
        scan_layout.addLayout(status_code_layout)
        
        # 请求方法选择
        method_layout = QHBoxLayout()
        method_layout.addWidget(QLabel("请求方法:"))
        self.method_combo = QComboBox()
        self.method_combo.addItems(["HEAD", "GET"])
        method_layout.addWidget(self.method_combo)
        scan_layout.addLayout(method_layout)
        
        scan_layout.addStretch()
        scan_group.setLayout(scan_layout)
        config_splitter.addWidget(scan_group)
        
        config_splitter.setSizes([600, 300])
        main_layout.addWidget(config_splitter)
        
        results_splitter = QSplitter(Qt.Vertical)
        
        # 路径组件显示
        components_group = QGroupBox("提取的路径组件")
        components_layout = QVBoxLayout()
        self.components_display = QTextEdit()
        self.components_display.setReadOnly(True)
        self.components_display.setLineWrapMode(QTextEdit.WidgetWidth)
        components_layout.addWidget(self.components_display)
        components_group.setLayout(components_layout)
        results_splitter.addWidget(components_group)
        
        # 扫描结果显示
        results_group = QGroupBox("扫描结果")
        results_layout = QVBoxLayout()
        self.result_output = QTextEdit()
        self.result_output.setReadOnly(True)
        self.result_output.setLineWrapMode(QTextEdit.WidgetWidth)
        results_layout.addWidget(self.result_output)
        results_group.setLayout(results_layout)
        results_splitter.addWidget(results_group)
        
        results_splitter.setSizes([300, 400])
        main_layout.addWidget(results_splitter)
        
        status_layout = QHBoxLayout()
        self.progress = QProgressBar()
        self.status_label = QLabel("准备就绪")
        self.current_url_label = QLabel("当前URL: 无")
        self.current_url_label.setMinimumWidth(300)
        self.current_url_label.setMaximumWidth(500)
        self.current_url_label.setStyleSheet("text-overflow: ellipsis; white-space: nowrap; overflow: hidden;")
        self.current_url_label.setTextInteractionFlags(Qt.TextSelectableByMouse)
        
        status_layout.addWidget(self.progress, 3)
        status_layout.addWidget(self.current_url_label, 5)
        status_layout.addWidget(self.status_label, 2)
        main_layout.addLayout(status_layout)
        
        self.setCentralWidget(central_widget)
        self.resize(1200, 800)
        
        self.setMinimumSize(1000, 600)

    def update_monitor(self):
        cpu_usage = SystemMonitor.get_cpu_usage()
        mem_usage = SystemMonitor.get_memory_usage()
        disk_usage = SystemMonitor.get_disk_usage()
        
        self.cpu_label.setText(f"CPU: {cpu_usage}%")
        self.mem_label.setText(f"内存: {mem_usage}")
        self.disk_label.setText(f"磁盘: {disk_usage}")

    def batch_import(self):
        """批量导入目标URL"""
        file_path, _ = QFileDialog.getOpenFileName(self, "选择目标文件", "", "文本文件 (*.txt)")
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    targets = [line.strip() for line in f if line.strip()]
                
                if not targets:
                    QMessageBox.warning(self, "警告", "文件中未找到有效目标")
                    return
                
                # 验证并筛选存活目标
                self.status_label.setText("正在验证目标有效性...")
                QApplication.processEvents()
                
                alive_targets = []
                for target in targets:
                    if TargetValidator.is_valid_target(target):
                        if TargetValidator.is_alive(target):
                            alive_targets.append(target)
                            self.status_label.setText(f"已验证 {len(alive_targets)}/{len(targets)} 个目标")
                            QApplication.processEvents()
                
                self.alive_targets = alive_targets
                self.current_target_index = 0
                self.batch_scanning = True
                self.batch_status.setText(f"批量扫描: 已启用 ({len(alive_targets)}个有效目标)")
                
                if alive_targets:
                    self.url_input.setText(alive_targets[0])
                    self.status_label.setText(f"已导入 {len(alive_targets)} 个存活目标")
                else:
                    self.batch_status.setText("批量扫描: 未启用")
                    QMessageBox.warning(self, "警告", "未找到存活的有效目标")
                    
            except Exception as e:
                QMessageBox.critical(self, "错误", f"导入失败: {str(e)}")

    def import_custom_dict(self):
        """导入自定义路径字典"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择字典文件", "", "文本文件 (*.txt);;所有文件 (*)"
        )
        if not file_path:
            return

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                self.custom_components = list({
                    line.strip() for line in f 
                    if line.strip() and not line.startswith('#')
                })
            
            if not self.custom_components:
                QMessageBox.warning(self, "警告", "字典文件为空或仅包含注释")
                return

            # 显示前10行预览
            preview_text = "\n".join(self.custom_components[:10])
            if len(self.custom_components) > 10:
                preview_text += "\n..."
            self.dict_preview.setPlainText(preview_text)
            self.dict_path_edit.setText(file_path)
            self.status_label.setText(f"已导入字典: {len(self.custom_components)} 条路径")
            self.fetch_btn.setEnabled(False)
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"导入字典失败: {str(e)}")

    def on_config_changed(self):
        """配置变更时延迟应用（防抖）"""
        if self.worker and self.worker.isRunning() and not self.is_paused:
            self.config_timer.start(500)

    def apply_config_changes(self):
        """应用配置变更（实时生效）"""
        if self.worker and self.worker.isRunning() and not self.is_paused:
            new_depth = self.depth_input.value()
            new_threads = self.thread_input.value()
            new_base_path = self.baseurl_input.text()
            new_delay = self.delay_input.value()
            new_retry = self.retry_input.value()
            
            # 仅在配置实际变更时发送信号
            if (new_depth != self.last_depth or 
                new_threads != self.last_threads or 
                new_base_path != self.last_basepath or
                new_delay != self.last_delay or
                new_retry != self.last_retry):
                
                self.worker.config_changed.emit(new_depth, new_threads, new_base_path, new_delay, new_retry)
                # 更新最后配置缓存
                self.last_depth = new_depth
                self.last_threads = new_threads
                self.last_basepath = new_base_path
                self.last_delay = new_delay
                self.last_retry = new_retry

    def get_filter_status_codes(self):
        """解析用户输入的状态码列表"""
        text = self.status_codes_edit.text().strip()
        if not text:
            return []
        try:
            return [int(code.strip()) for code in text.split(',') if code.strip().isdigit()]
        except ValueError:
            QMessageBox.warning(self, "警告", "状态码格式错误，请输入数字并用逗号分隔")
            return [200, 403, 301, 302]

    def fetch_urls(self):
        """提取URL路径组件"""
        target_url = self.url_input.text().strip()
        if not target_url:
            QMessageBox.warning(self, "警告", "请输入目标URL")
            return

        self.status_label.setText("正在从JS文件中提取URL...")
        QApplication.processEvents()

        try:
            urls = JSFinder.find_urls(target_url)
            self.components = PathProcessor.extract_components(urls)
            self.components_display.setPlainText('\n'.join(self.components))
            self.status_label.setText(f"已提取 {len(self.components)} 个路径组件")
            self.start_btn.setEnabled(True)
            # 清除自定义字典
            if hasattr(self, 'custom_components'):
                delattr(self, 'custom_components')
                self.dict_preview.clear()
                self.dict_path_edit.clear()
            self.import_dict_btn.setEnabled(True)
            
        except Exception as e:
            self.status_label.setText(f"提取URL失败: {str(e)}")
            QMessageBox.critical(self, "错误", f"提取URL失败: {str(e)}")

    def start_scan(self):
        """开始扫描（支持自定义字典）"""
        # 优先使用自定义字典
        if hasattr(self, 'custom_components') and self.custom_components:
            components = self.custom_components
        elif hasattr(self, 'components') and self.components:
            components = self.components
        else:
            QMessageBox.warning(self, "警告", "请先提取路径组件或导入自定义字典")
            return

        base_url = self.url_input.text().strip()
        if not base_url:
            QMessageBox.warning(self, "警告", "请输入目标URL")
            return

        max_depth = self.depth_input.value()
        max_threads = self.thread_input.value()
        base_url_path = self.baseurl_input.text()
        request_delay = self.delay_input.value()
        retry_count = self.retry_input.value()

        self.result_output.clear()
        self.progress.setValue(0)
        self.start_btn.setEnabled(False)
        self.pause_btn.setEnabled(True)
        self.stop_btn.setEnabled(True)
        self.export_btn.setEnabled(False)
        self.status_label.setText("扫描中...")
        self.import_dict_btn.setEnabled(False)
        self.fetch_btn.setEnabled(False)
        self.is_paused = False
        
        # 创建当前扫描的临时文件
        timestamp = time.strftime('%Y%m%d_%H%M%S')
        target_name = re.sub(r'[^\w\-_.]', '_', base_url)[:50]
        self.current_scan_file = os.path.join(self.results_dir, f"scan_{target_name}_{timestamp}.txt")
        
        # 保存当前配置作为最后使用的配置
        self.last_depth = max_depth
        self.last_threads = max_threads
        self.last_basepath = base_url_path
        self.last_delay = request_delay
        self.last_retry = retry_count

        # 创建工作线程
        self.worker = BruteWorker(
            base_url, 
            components, 
            max_depth, 
            max_threads, 
            base_url_path,
            request_delay,
            retry_count
        )
        self.worker.update_progress.connect(self.update_progress)
        self.worker.found_path.connect(self.add_result)
        self.worker.error_path.connect(self.add_error)
        self.worker.status_update.connect(self.status_label.setText)
        self.worker.finished.connect(self.scan_finished)
        self.worker.paused.connect(self.on_paused_state_changed)
        self.worker.thread_status.connect(self.update_thread_count)
        self.worker.current_url.connect(self.update_current_url)
        self.worker.config_updated.connect(self.status_label.setText)
        self.worker.start()

    def toggle_pause(self):
        if self.worker and self.worker.isRunning():
            with self.worker.lock:
                current_pause = self.worker.pause_flag
            
            if current_pause:
                self.worker.resume()
            else:
                self.worker.pause()

    def on_paused_state_changed(self, is_paused):
        self.pause_btn.setText("继续" if is_paused else "暂停")
        self.is_paused = is_paused
        if is_paused:
            self.import_dict_btn.setEnabled(True)
            self.fetch_btn.setEnabled(True)
        else:
            self.import_dict_btn.setEnabled(False)
            self.fetch_btn.setEnabled(False)

    def stop_scan(self):
        if self.worker and self.worker.isRunning():
            self.is_stopping = True
            self.worker.stop()
            self.status_label.setText("正在停止...")
            self.start_btn.setEnabled(False)
            self.pause_btn.setEnabled(False)
            self.pause_btn.setText("暂停")
            self.import_dict_btn.setEnabled(True)
            self.fetch_btn.setEnabled(True)

    def update_progress(self, progress, completed):
        """更新进度条"""
        self.progress.setValue(progress)

    def update_thread_count(self, count):
        """更新线程计数显示"""
        self.thread_label.setText(f"线程: {count}")

    def update_current_url(self, url):
        """更新当前扫描的URL（长URL处理）"""
        # 限制URL显示长度，超出部分自动省略
        max_display_length = 80
        if len(url) > max_display_length:
            display_url = f"{url[:40]}...{url[-35:]}"
        else:
            display_url = url
        self.current_url_label.setText(f"当前URL: {display_url}")

    def add_result(self, url, info, depth):
        """添加成功结果（支持状态码过滤和实时保存）"""
        # 从info中提取状态码
        status_code_match = re.search(r'状态码: (\d+)', info)
        if not status_code_match:
            return

        status_code = int(status_code_match.group(1))
        filter_codes = self.get_filter_status_codes()

        # 只显示过滤后的状态码
        if not filter_codes or status_code in filter_codes:
            result_text = f"✅ 深度 {depth} - {url} [{info}]"
            self.result_output.append(result_text)
            self.success_paths.append(result_text)
            
            # 实时保存到文件
            if self.current_scan_file:
                try:
                    with open(self.current_scan_file, 'a', encoding='utf-8') as f:
                        f.write(result_text + "\n")
                except Exception as e:
                    print(f"保存结果失败: {str(e)}")

    def add_error(self, url, msg):
        """添加错误结果"""
        self.result_output.append(f"❌ {url} - {msg}")

    def scan_finished(self):
        """扫描完成处理"""
        if self.batch_scanning and not self.is_stopping:
            self.current_target_index += 1
            if self.current_target_index < len(self.alive_targets):
                next_target = self.alive_targets[self.current_target_index]
                self.url_input.setText(next_target)
                self.status_label.setText(f"准备扫描下一个目标: {next_target}")
                
                # 自动提取下一个目标的路径组件并开始扫描
                QTimer.singleShot(1000, self.auto_next_scan)
                return
        
        # 非批量扫描或已完成所有批量目标
        self.worker = None
        self.start_btn.setEnabled(True)
        self.pause_btn.setEnabled(False)
        self.stop_btn.setEnabled(False)
        self.export_btn.setEnabled(bool(self.success_paths))
        self.thread_label.setText("线程: 0")
        self.current_url_label.setText("当前URL: 无")
        self.is_stopping = False
        self.pause_btn.setText("暂停")
        self.import_dict_btn.setEnabled(True)
        self.fetch_btn.setEnabled(True)
        self.is_paused = False
        
        if self.batch_scanning and self.current_target_index >= len(self.alive_targets):
            self.status_label.setText("所有批量目标扫描完成")
            self.batch_scanning = False
            self.batch_status.setText("批量扫描: 未启用")

    def auto_next_scan(self):
        try:
            urls = JSFinder.find_urls(self.url_input.text().strip())
            self.components = PathProcessor.extract_components(urls)
            self.components_display.setPlainText('\n'.join(self.components))
            self.status_label.setText(f"已提取 {len(self.components)} 个路径组件，准备开始扫描")
            QTimer.singleShot(1000, self.start_scan)
        except Exception as e:
            self.status_label.setText(f"提取URL失败: {str(e)}")
            self.batch_scanning = False
            self.batch_status.setText("批量扫描: 未启用")
            self.start_btn.setEnabled(True)
            self.stop_btn.setEnabled(False)
            self.import_dict_btn.setEnabled(True)
            self.fetch_btn.setEnabled(True)

    def export_results(self):
        if not self.success_paths:
            self.status_label.setText("没有可导出的结果")
            return

        file_name = f"all_scan_results_{time.strftime('%Y%m%d_%H%M%S')}.txt"
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存所有结果", os.path.join(self.results_dir, file_name), "文本文件 (*.txt)"
        )
        
        if file_path:
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write("\n".join(self.success_paths))
                self.status_label.setText(f"所有结果已导出到 {file_path}")
            except Exception as e:
                self.status_label.setText(f"导出失败: {str(e)}")
                QMessageBox.critical(self, "错误", f"导出失败: {str(e)}")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
